import pygame
import numpy as np

class Terrain:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.tile_size = 32  # 每个格子的像素大小
        self.tile_height = 16  # 每个格子的高度
        
        # 创建地形高度图
        self.height_map = np.zeros((width, height))
        
        # 计算地形网格的偏移量
        self.offset_x = 400  # 屏幕中心x坐标
        self.offset_y = 100  # 屏幕顶部y坐标
    
    def screen_to_grid(self, screen_pos):
        # 将屏幕坐标转换为网格坐标
        x, y = screen_pos
        x -= self.offset_x
        y -= self.offset_y
        
        # 斜45度视角的坐标转换
        grid_x = int((x / self.tile_size + y / self.tile_height) / 2)
        grid_y = int((y / self.tile_height - x / self.tile_size) / 2)
        
        # 确保坐标在有效范围内
        if 0 <= grid_x < self.width and 0 <= grid_y < self.height:
            return (grid_x, grid_y)
        return None
    
    def grid_to_screen(self, grid_x, grid_y):
        # 将网格坐标转换为屏幕坐标
        screen_x = (grid_x - grid_y) * self.tile_size + self.offset_x
        screen_y = (grid_x + grid_y) * self.tile_height + self.offset_y
        return (screen_x, screen_y)
    
    def draw(self, screen):
        # 绘制地形网格
        for x in range(self.width):
            for y in range(self.height):
                screen_x, screen_y = self.grid_to_screen(x, y)
                
                # 绘制地形格子
                points = [
                    (screen_x, screen_y),
                    (screen_x + self.tile_size, screen_y - self.tile_height),
                    (screen_x + self.tile_size * 2, screen_y),
                    (screen_x + self.tile_size, screen_y + self.tile_height)
                ]
                
                # 根据高度设置不同的颜色
                height = self.height_map[x, y]
                color = (50 + height * 20, 50 + height * 20, 50 + height * 20)
                pygame.draw.polygon(screen, color, points)
                pygame.draw.polygon(screen, (100, 100, 100), points, 1)  # 绘制边框